import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

#构造模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 初始化网络结构
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.bn1 = nn.BatchNorm2d(20)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.bn2 = nn.BatchNorm2d(50)
        self.fc1 = nn.Linear(4 * 4 * 50, 500)
        self.bn3 = nn.BatchNorm1d(500)
        self.fc2 = nn.Linear(500, 10)
        self.dropout = nn.Dropout(0.5)  # 添加Dropout层

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.max_pool2d(x, 2, 2)
        x = x.view(-1, 4 * 4 * 50)
        x = F.relu(self.bn3(self.fc1(x)))
        x = self.dropout(x)  # 在全连接层之前添加Dropout
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)
# 训练集数据预处理和加载
def setup_train_data_loader(batch_size, rank, world_size):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)
    return dataloader
# 多进程通信环境初始化
def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    #在多个进程之间建立通信的桥梁，使得它们可以相互发送和接收数据，协作完成训练任务
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    torch.manual_seed(0)
# 清理过程
def cleanup():
    dist.destroy_process_group()

'''循环训练：rank代表进程ID，world_size代表进程总数'''
def train(rank, world_size):
    #多进程通信环境初始化
    setup(rank, world_size)
    #训练数据初始化
    batch_size = 64
    dataloader = setup_train_data_loader(batch_size, rank, world_size)
    '''torch.device() 函数用于指定用于计算的设备
    0 是指定使用第一个GPU。如果你的系统上有多个GPU，你可以通过更改 0 为 1、2 等来选择不同的GPU
    '''
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    #print(f'device：{device}')
    #移动模型：将自定义的Net模型的参数移动到之前设置的device设备上，以便在该设备上进行训练和推理
    model = Net().to(device)
    '''
    将模型复制到多个GPU上，并进行数据并行处理，从而加快训练速度。（只不过我的电脑上只有一个GPU）
    model：这是之前创建的神经网络模型实例，它包含了模型的结构和参数
    第二个参数使用device_ids=[0] if torch.cuda.is_available() else None条件表达式来决定放在GPU上运行还是CPU上运行
    '''
    ddp_model = DDP(model, device_ids=[0] if torch.cuda.is_available() else None)
    '''将损失函数实例移动到device设备上（也就是第一个GPU上）'''
    criterion = nn.CrossEntropyLoss().to(device)
    '''修改思路：使用自定义的优化器；然后修改梯度下降的部分'''
    optimizer = optim.Adam(ddp_model.parameters(), lr=0.001)
    #通过设置range()中的参数来设置训练的轮数
    for epoch in range(1):
        ddp_model.train() #model.train()
        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = ddp_model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            #打印每个进程训练自己拥有的训练集的loss
            if batch_idx % 10 == 0:
                print(f'Rank {rank}, Epoch {epoch}, Batch {batch_idx}, Loss {loss.item()}')       
    # 保存训练好的模型
    torch.save(model.state_dict(), 'model.pth')
    cleanup()
    print("训练结束")
# 测试函数
def test(model_path):
    # 测试数据集的加载
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)
    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

    # 加载模型
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = Net().to(device)
    model.load_state_dict(torch.load(model_path))  # 加载保存的模型状态
    model.eval()  # 设置模型为评估模式
    # 初始化测试损失
    test_loss = 0.0
    # 测试模型
    with torch.no_grad():  # 关闭梯度计算，节省计算资源
        for batch_idx, (data, target) in enumerate(test_loader):
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = F.cross_entropy(output, target)

            # 累计损失
            test_loss += loss.item() * data.size(0)

    # 计算平均测试损失
    test_loss /= len(test_loader.dataset)

    # 打印最终的测试损失
    print(f'Average Test Loss: {test_loss}')

# 启动训练
def main():
    world_size = 5  # 模拟3个进程
    '''（分布式训练关键）mp.spawn函数是创建多个进程来模拟多GPU训练神经网络模型
    mp是torch.multiprocessing 
    参数：
    据指定的nprocs（即进程数）创建多个进程
    args是给train函数的参数值
    train：想要在每个进程中执行的函数。这个函数通常包含训练循环和模型训练的逻辑。
    join=True：这个参数表示是否等待所有子进程完成
    整体效果：
    在分布式训练中，每个进程都会处理一部分数据，并且会在每个epoch结束后进行参数同步，以确保所有进程看到的模型参数是一致的。
    '''
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)
    test('model.pth')
if __name__ == '__main__':
    main()
